# 基础 RNN Cell 使用
import torch

print('===RNN Cell的使用===')
# 定义batch大小、序列长度、输入样本特征维度、输出维度大小
batch_size = 1
seq_len = 3
input_size = 4
hidden_size = 2

# 定义 RNN 单元，这里简单地使用一个RNN单元来模拟
cell = torch.nn.RNNCell(input_size=input_size, hidden_size=hidden_size)
# 简单定义一个数据集，维度为 (seq, batch, features)
dataset = torch.randn(seq_len, batch_size, input_size)
# 初始化一个隐层输入h0，因为t0时刻是没有隐层输入的
hidden = torch.zeros(batch_size, hidden_size)

# 训练，观察输入和隐层输出维度的变化
for idx, inputs in enumerate(dataset):
    print('=' * 20, idx, '=' * 20)
    print('Input size: ', inputs.shape)
    # 得到隐层输出
    hidden = cell(inputs, hidden)
    print('outputs size: ', hidden.shape)
    print(hidden)

# 基础RNN的使用

print('===RNN的使用===')
batch_size = 1
seq_len = 3
input_size = 4
hidden_size = 2
# RNN 结构中隐层的层数
num_layers = 1
# RNN 网络结构定义
cell = torch.nn.RNN(input_size=input_size, hidden_size=hidden_size,
                    num_layers=num_layers)
# (seqLen, batchSize, inputSize)
inputs = torch.randn(seq_len, batch_size, input_size)
# 初始化t0时刻隐层输入
hidden = torch.zeros(num_layers, batch_size, hidden_size)
# 训练
out, hidden = cell(inputs, hidden)
# 比较输入输出的维度
print('Output size:', out.shape)
print('Output:', out)
print('Hidden size: ', hidden.shape)
print('Hidden: ', hidden)
